Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@src/megatron/bridge/recipes/qwen_vl/data/energon/task_encoder.py`:
- Around line 174-179: The code calls pickle.loads(data) inside __call__ with
data from webdataset shards, which is unsafe for untrusted inputs; update the
function to either (A) document the trust assumption clearly (add a comment
above __call__ stating that data must come from trusted/curated sources), or (B)
replace the deserialization with a safer format/loader (e.g., JSON/msgpack or a
vetted safe_unpickle helper) and handle errors with a try/except to avoid
crashing on malformed data; reference the __call__ method and the extension
check (self.extensions) when making the change so the logic and error handling
remain consistent.
- Around line 188-205: Docstring claims custom handlers for image, audio, and
video, but ChatMLWebdataset.__init__ only registers imagehandler and
videohandler on self._decoder; update the implementation or docstring: either
register the missing audio handler (e.g., add audiohandler(self.audio_decode) to
the Decoder list) and ensure an appropriate audio_decode method exists, or
change the class docstring to remove "audio" so it accurately reflects handlers
currently registered (imagehandler and videohandler).
- Line 16: Remove the unused import causing the linter failure: delete the
top-level "import io" statement in
src/megatron/bridge/recipes/qwen_vl/data/energon/task_encoder.py so no unused
modules remain; ensure no other references to "io" exist in functions or classes
within that module (e.g., TaskEncoder or any helper functions) before
committing.

---

Nitpick comments:
In `@src/megatron/bridge/data/energon/base_energon_datamodule.py`:
- Around line 187-194: cp_size and cp_rank are only retrieved for observability
but not used elsewhere; add a short inline comment next to the calls to
parallel_state.get_context_parallel_world_size() and
parallel_state.get_context_parallel_rank() (and/or above the logger.info block)
stating that these variables are logged for debugging/observability and
intentionally not included in WorkerConfig (e.g., "# Logged for debugging; not
used in WorkerConfig") so future readers know their purpose.

In `@src/megatron/bridge/recipes/qwen_vl/data/energon/task_encoder.py`:
- Around line 140-154: _in _resolve_hf_mm_token_ids, tighten the except clause
in helper _get so we don't swallow unexpected errors; replace the broad "except
Exception" around hf_tokenizer.convert_tokens_to_ids(token_str) with a narrower
catch like "except (KeyError, ValueError, TypeError)" so only anticipated
token-conversion failures fall back to the default_id while other bugs (e.g.,
AttributeError or programming errors) still surface; update both calls that
compute image_id and video_id via _get accordingly.
- Line 23: Replace typing generics with Python 3.10+ built-ins and use | None
for nullable types: remove imports of Dict, List, Optional and update all
occurrences of Dict[...] -> dict[...], List[...] -> list[...], and Optional[T]
-> T | None in this module (referencing the symbols Dict, List, Optional in the
import and any function signatures/annotations inside task_encoder.py such as
encode_task or other typed variables). Ensure all type hints compile under
Python 3.10+ and adjust any forward references if needed.
- Around line 166-185: Rename the class videohandler to PascalCase VideoHandler
and fix its docstring to "Create a video handler"; update all references to the
class (e.g., constructor usages and imports such as ChatMLWebdataset.__init__,
any test imports and test usage that reference videohandler) so they refer to
VideoHandler to avoid breakage, ensuring the class name change is consistent
across the module, initializations, and tests.

In `@tests/functional_tests/data/energon/test_base_energon_datamodule.py`:
- Around line 177-183: These tests are pure unit tests but live in
tests/functional_tests; move the mocked-only test classes
TestEnergonDataModuleCPHandling and TestEnergonDataShardingVerification into
tests/unit_tests/, update any imports or test fixtures they rely on, and ensure
their pytest markers (if any) still apply; keep the test code unchanged other
than adjusting module/package imports and test path so CI treats them as unit
tests.

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@examples/models/vlm/qwen3_vl/energon_test.sh`:
- Line 2: Update the copyright header string at the top of energon_test.sh which
currently reads "Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved."
to use the correct year 2026; locate the header line in the file and change the
year token from 2025 to 2026 so the file reflects the PR creation year.
- Around line 1-17: The script lacks strict shell options; insert "set -euo
pipefail" immediately after the shebang in energon_test.sh to ensure the script
exits on errors, treats unset variables as failures, and propagates pipe
failures; update startup of the script (around the top where WORKSPACE is
defined) so the shell options are set before any variable expansion or commands
execute.
- Line 49: Replace the direct Python invocation with the project's wrapper:
change the command that currently starts with "python -m torch.distributed.run
--nproc_per_node=$N_PROC scripts/training/run_recipe.py" to use "uv run" (e.g.,
"uv run python -m torch.distributed.run --nproc_per_node=$N_PROC
scripts/training/run_recipe.py") so the script invocation for
torch.distributed.run / scripts/training/run_recipe.py follows the coding
guideline requiring uv run in shell/example scripts.
- Line 48: The echo currently prints DP=$N_PROC which is wrong because N_PROC is
total processes; compute actual data-parallel degree DP as DP=$(( N_PROC / (EP *
TP * PP * CP) )) (or set DP earlier when you compute it) and update the echo
line to print DP (and/or include the formula) instead of N_PROC; refer to the
variables N_PROC, EP, TP, PP, CP and the echo statement that currently contains
"DP=$N_PROC" to locate and fix the line.

In `@src/megatron/bridge/recipes/qwen_vl/data/energon/task_encoder.py`:
- Around line 165-166: Rename the class `videohandler` to `VideoHandler` and
update its docstring from "Create an video handler." to "Create a video
handler."; then update all references/imports that use `videohandler` (e.g.,
test imports like `from ... import videohandler`) and any `dataset.yaml` or
config entries that reference the class by name so they use `VideoHandler`
instead to keep naming consistent with the coding guidelines.
- Line 179: The code checks extension.lower() for membership but then uses the
original-cased variable `extension` to index `self.extensions_mapping`, which
causes KeyError for mixed/upper-case inputs; update the code to normalize the
extension before lookup (e.g., compute a lowercased `ext_lower =
extension.lower()` or reassign `extension = extension.lower()`), use that
normalized value for both the guard and when computing `key =
self.extensions_mapping[...]`, and ensure any subsequent usage in the same scope
uses the normalized name so lookups against `self.extensions_mapping` succeed.

---

Outside diff comments:
In `@src/megatron/bridge/data/energon/energon_provider.py`:
- Around line 36-49: The pack_sequences_in_batch flag on EnergonProvider is
never forwarded in build_datasets, so setting pack_sequences_in_batch=True is a
no-op; update build_datasets to pass self.pack_sequences_in_batch into the
EnergonMultiModalDataModule constructor (or, if EnergonMultiModalDataModule does
not accept this parameter yet, detect self.pack_sequences_in_batch and raise
NotImplementedError to fail fast). Locate the flag on the EnergonProvider class,
the build_datasets method, and the EnergonMultiModalDataModule constructor to
wire the parameter or add the explicit error handling.

---

Duplicate comments:
In `@src/megatron/bridge/recipes/qwen_vl/data/energon/task_encoder.py`:
- Around line 173-178: In __call__, do not use pickle.loads directly because of
unsafe deserialization; instead either decode JSON (e.g.,
json.loads(data.decode())) if the payload is JSON, or, if you must unpickle,
replace pickle.loads(data) with a safe unpickler: implement a
RestrictedUnpickler that overrides find_class and only allows a whitelist of
expected classes/modules, then call
RestrictedUnpickler(io.BytesIO(data)).load(); update the variable 'data'
assignment in __call__ to use that safe method.

---

Nitpick comments:
In `@src/megatron/bridge/recipes/qwen_vl/data/energon/task_encoder.py`:
- Line 22: Replace typing-based generics and Optional usage with built-in
generics and PEP 604 union syntax: remove Dict, List, Optional imports and
update annotations in the ChatMLSample fields (the new fields around the
ChatMLSample class at lines ~161–162) to use dict[...] and list[...] and use T |
None instead of Optional[T]; ensure any other occurrences in this module follow
the same change (e.g., function signatures or field types referencing
Dict/List/Optional) and run type checks to confirm no remaining typing aliases
are used.